package router

import (
	"fmt"
	"github.com/tmc/langchaingo/outputparser"
	"github.com/tmc/langchaingo/prompts"
	"strings"
)

const (
	_formatting   = "formatting"
	_destinations = "destinations"
	_input        = "input"
	_nextInput    = "next_inputs"

	MULTI_PROMPT_ROUTER_TEMPLATE = `Given a raw text input to a language model select the model prompt best suited for 
the input. You will be given the names of the available prompts and a description of 
what the prompt is best suited for. You may also revise the original input if you 
think that revising it will ultimately lead to a better response from the language 
model.

<< FORMATTING >>
Return a markdown code snippet with a JSON object formatted to look like:
{{.formatting}}

REMEMBER: "destination" MUST be one of the candidate prompt names specified below OR 
it can be "DEFAULT" if the input is not well suited for any of the candidate prompts.
REMEMBER: "next_inputs" can just be the original input if you don't think any 
modifications are needed.

<< CANDIDATE PROMPTS >>
{{.destinations}}

<< INPUT >>
{{.input}}
`
)

var (
	_outputparser = outputparser.NewStructured([]outputparser.ResponseSchema{
		{
			Name:        _destinations,
			Description: `name of the question answering system to use or "DEFAULT"`,
		}, {
			Name:        _nextInput,
			Description: `a potentially modified version of the original input`,
		},
	})
)

func createPrompt(handler []Handler) prompts.PromptTemplate {
	return prompts.PromptTemplate{
		Template:       MULTI_PROMPT_ROUTER_TEMPLATE,
		InputVariables: []string{_input},
		TemplateFormat: prompts.TemplateFormatGoTemplate,
		PartialVariables: map[string]any{
			//"destinations":
			_destinations: handlerDestinations(handler),
			_formatting:   _outputparser.GetFormatInstructions(),
		},
	}
}

func handlerDestinations(handler []Handler) string {
	var hs strings.Builder
	for _, h := range handler {
		hs.WriteString(fmt.Sprintf("- %s: %s\n", h.Name(), h.Description()))
	}

	return hs.String()
}
